/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: GPU/NPU send/recv window memory manager
 */
#include "typical_window_mem.h"
#include "adapter_rts_common.h"

namespace hccl {
TypicalWindowMem &TypicalWindowMem::GetInstance()
{
    static TypicalWindowMem typicalWindowMem;
    return typicalWindowMem;
}

TypicalWindowMem::TypicalWindowMem()
{
}

TypicalWindowMem::~TypicalWindowMem()
{
    FreeAllWinowMem();
}

HcclResult TypicalWindowMem::AllocWindowMem(void **ptr, uint64_t size)
{
    HCCL_DEBUG("[TypicalWindowMem][AllocWindowMem]start alloc window mem to [%p] of size[%llu].", ptr, size);
    // Check ptr and size validation.
    CHK_PTR_NULL(ptr);
    CHK_PRT_RET(size < 1,
        HCCL_ERROR("[TypicalWindowMem][AllocWindowMem]size[%lu], cannot alloc window mem less than 1.", size),
        HCCL_E_PARA);

    // Check whether the addr already exist in window mem map.
    std::unique_lock<std::mutex> lockWindowMemMap(windowMemMapMutex_);
    if (*ptr != nullptr) {
        uint64_t inputAddr = reinterpret_cast<uintptr_t>(*ptr);
        auto wmIter = windowMemMap_.find(inputAddr);
        if (wmIter != windowMemMap_.end()) {
            HCCL_ERROR("[TypicalWindowMem][AllocWindowMem]addr[%p] already allocated.", *ptr);
            return HCCL_E_PARA;
        }
        HCCL_WARNING("[TypicalWindowMem][AllocWindowMem]addr[%p] is not nullptr, " \
            "we will overwrite it with the new allocated address.", *ptr);
    }
    lockWindowMemMap.unlock();

    // Alloc window mem.
    CHK_RET(hrtMalloc(ptr, size));
    CHK_PRT_RET((*ptr) == nullptr,
        HCCL_ERROR("[TypicalWindowMem][AllocWindowMem]In typical notify src buffer, malloc failed."), HCCL_E_MEMORY);

    // Add allocated mem into window mem map。
    uint64_t allocatedAddr = reinterpret_cast<uintptr_t>(*ptr);
    lockWindowMemMap.lock();
    windowMemMap_[allocatedAddr] = size;
    lockWindowMemMap.unlock();

    HCCL_INFO("[TypicalWindowMem][AllocWindowMem]alloc window memory success, addr[%p], size[%llu]. " \
        "please register mr before use", *ptr, size);
    return HCCL_SUCCESS;
}

HcclResult TypicalWindowMem::FreeWindowMem(void *ptr)
{
    HCCL_DEBUG("[TypicalWindowMem][FreeWindowMem]start free window mem on [%p], please deregister mr before free.",
        ptr);
    CHK_PTR_NULL(ptr);
    // Check whether the addr exist in window mem map. Remove allocated mem from window mem map if exists.
    uint64_t addr = reinterpret_cast<uintptr_t>(ptr);
    std::unique_lock<std::mutex> lockWindowMemMap(windowMemMapMutex_);
    auto wmIter = windowMemMap_.find(addr);
    if (wmIter == windowMemMap_.end()) {
        HCCL_ERROR("[TypicalWindowMem][AllocWindowMem]addr[%p] were not allocated or already freed.", ptr);
        return HCCL_E_PARA;
    }
    lockWindowMemMap.unlock();

    // Free window mem.
    CHK_RET(hrtFree(ptr));
    
    lockWindowMemMap.lock();
    windowMemMap_.erase(wmIter);
    lockWindowMemMap.unlock();
    ptr = nullptr;
    HCCL_INFO("[TypicalWindowMem][AllocWindowMem]free window memory success, addr[%p].", ptr);
    return HCCL_SUCCESS;
}

HcclResult TypicalWindowMem::FreeAllWinowMem()
{
    std::unique_lock<std::mutex> lockWindowMemMap(windowMemMapMutex_);
    if (!windowMemMap_.empty()) {
        for (auto &wmIter : windowMemMap_) {
            auto ptr = reinterpret_cast<void *>(static_cast<uintptr_t>(wmIter.first));
            if (ptr != nullptr) {
                CHK_RET(hrtFree(ptr));
            }
        }
        windowMemMap_.clear();
    }
    HCCL_INFO("[TypicalWindowMem][FreeAllWinowMem]free all window memory success.");
    return HCCL_SUCCESS;
}
}   // namespace hccl
